/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package main.NER.Trie;

import java.io.Serializable;
import java.util.*;
import java.util.regex.Pattern;
import main.ContextFree.DynamicRegexGenerator;
import main.EntityClassGraph.ElementOntologySetGraph;
import main.KBProcess.Element;
import main.KBProcess.Fragment;
import main.KBProcess.KnowledgeBase;
import main.KBProcess.KnowledgeBaseDocument;
import main.NER.Trie.ContextTrie.ORIENTATION;
import main.Ontology.ElementOntologyFactory;
import main.Ontology.ElementOntologySet;
import main.ContextFree.RegexCompare;

/**
 * KB Trie stores the Trie data structure for both left and right <code>ContextTrie
 * </code>. It is the object representation of all the information extracted from the
 * annotated input data
 * @author Isaac Osesina
 */
public class KBTrie implements Serializable, Cloneable  {

    List<KnowledgeBase> kbList;
    private ContextTrie LEFT_CONTEXT_TRIE;
    private ContextTrie RIGHT_CONTEXT_TRIE;
    public static double threshold = 0.5; //0.1
    private Map<ElementOntologySet, Pattern> allEntityTypeIDRegex;
    private RegexCompare patternComparison;
    public boolean normalize;
    private Map<Element, Map<ElementOntologySet, Pattern>> regexWithoutEntity;
    private ElementOntologyFactory elementOntologyFactory;
    private ElementOntologySetGraph elementOntologySetGraph;

    public KBTrie(KBTrie kbTrie) {
        kbList = kbTrie.kbList;
        LEFT_CONTEXT_TRIE = kbTrie.LEFT_CONTEXT_TRIE;
        RIGHT_CONTEXT_TRIE = kbTrie.RIGHT_CONTEXT_TRIE;
        allEntityTypeIDRegex = kbTrie.allEntityTypeIDRegex;
        patternComparison = kbTrie.patternComparison;
        normalize = kbTrie.normalize;
        regexWithoutEntity = kbTrie.regexWithoutEntity;
        elementOntologyFactory = kbTrie.elementOntologyFactory;
        elementOntologySetGraph = kbTrie.elementOntologySetGraph;
        
    }
    
    public KBTrie(ElementOntologyFactory elementOntologyFactory) {
        this.elementOntologyFactory = elementOntologyFactory;
    }
    

    public KBTrie(ElementOntologyFactory elementOntologyFactory, KnowledgeBase kb) {
        this.elementOntologyFactory = elementOntologyFactory;
        kbList = new ArrayList<KnowledgeBase>();
        kbList.add(kb);
        addKBToModel(kb);
    }

    /**
     * Returns left or right context trie
     * left context trie if orientation = ORIENTATION.LEFT
     * right context trie if orientation = ORIENTATION.RIGHT
     * @param orientation orientation of context trie
     * @return left or right <code>ContextTrie</code>
     */
    public ContextTrie getCONTEXT_TRIE(ORIENTATION orientation) {
        if (orientation == ORIENTATION.LEFT) {
            return LEFT_CONTEXT_TRIE;
        } else {
            return RIGHT_CONTEXT_TRIE;
        }
    }

    /**
     * Adds a <code>KnowledgeBase</code> to the KBTrie
     * @param kb KnowledgeBase to be added to KBTrie
     */
    public void addKB(KnowledgeBase kb) {
        kbList = kbList == null ? new ArrayList<KnowledgeBase>() : kbList;
        kbList.add(kb);
        addKBToModel(kb);
    }

    /**
     * Converts the <code>KnowledgeBase</code> into Trie data structure.
     * if KBTrie has been previously normalized, it would be re-normalized 
     * before being used to extract candidates
     * @param kb KnowledgeBase to be added to KBTrie
     */
    public void addKBToModel(KnowledgeBase kb) {
        LEFT_CONTEXT_TRIE = LEFT_CONTEXT_TRIE == null ? new ContextTrie(
                elementOntologyFactory, ORIENTATION.LEFT) : LEFT_CONTEXT_TRIE;
        RIGHT_CONTEXT_TRIE = RIGHT_CONTEXT_TRIE == null ? new ContextTrie(
                elementOntologyFactory, ORIENTATION.RIGHT) : RIGHT_CONTEXT_TRIE;

        for (Fragment f : kb.getFragmentList()) {
            addFragmentToModel(f);
        }
        normalize = false;
    }

    /**
     * Returns the left context trie
     * @return left context trie
     */
    public ContextTrie getLEFT_CONTEXT_TRIE() {
        return LEFT_CONTEXT_TRIE;
    }

    /**
     * Returns right context trie
     * @return right context trie
     */
    public ContextTrie getRIGHT_CONTEXT_TRIE() {
        return RIGHT_CONTEXT_TRIE;
    }    
    
    /**
     * Returns <code>RegexCompare</code>
     * This class contains information about the regular expression of all entity values
     * It is useful in determining entity types (ElementOntologySet) with overlapping
     * regular expression
     * @return entityClassRegexCompare
     */
    public RegexCompare getEntityClassRegexCompare() {
        return patternComparison;
    }
    
    /**
     * Returns a Map containing the regular expression of all entity types 
     * (ElementOntologySet) in this KBTrie
     * @return Map containing the regular expression of all entity types in this KBTrie
     */
    public Map<ElementOntologySet, Pattern> getAllEntityTypeIDRegex() {
        if(allEntityTypeIDRegex==null)
            return new HashMap<ElementOntologySet, Pattern>();
        return allEntityTypeIDRegex;
    }    
    
    
    /**
     * Adds <code>Fragment</code> to KBTrie
     * @param fragment to be added to KBTrie
     */
    public void addFragmentToModel(Fragment fragment) {
        LEFT_CONTEXT_TRIE = LEFT_CONTEXT_TRIE == null ? new ContextTrie(
                elementOntologyFactory, ORIENTATION.LEFT) : LEFT_CONTEXT_TRIE;
        LEFT_CONTEXT_TRIE.put(fragment);
        
        RIGHT_CONTEXT_TRIE = RIGHT_CONTEXT_TRIE == null ? new ContextTrie(
                elementOntologyFactory, ORIENTATION.RIGHT) : RIGHT_CONTEXT_TRIE;
        RIGHT_CONTEXT_TRIE.put(fragment);
        
        elementOntologySetGraph = elementOntologySetGraph == null ? 
                new ElementOntologySetGraph(): elementOntologySetGraph;
        elementOntologySetGraph.addVertex(fragment.getEntityElement());
        normalize = false;
    }
    

    /**
     * Returns a list of the <code>KnowledgeBase</code> contained in the KBTrie
     * @return a list of the KnowledgeBase in the KBTrie
     */
    public List<KnowledgeBase> getKbList() {
        return kbList;
    }

    /**
     * Returns a list of all the KB Fragments contained in the all the KnowledgeBase
     * in the KBTrie
     * @return list of all the KB Fragments
     */
    public List<Fragment> getFragmentList() {
        List<Fragment> fragmentList = null;
        if (kbList != null) {
            fragmentList = new ArrayList<Fragment>();
            for (KnowledgeBase kb : kbList) {
                if (kb != null) {
                    if (kb.getFragmentList() != null) {
                        fragmentList.addAll(kb.getFragmentList());
                    }
                }
            }
        }
        return fragmentList;
    }

    /**
     * Returns a list of all the <code>KnowledgeBaseDocument</code> contained in
     * the KBTrie
     * @return list of KnowledgeBaseDocument
     */
    public List<KnowledgeBaseDocument> getKBDocuments() {
        List<KnowledgeBaseDocument> fragmentList = null;
        if (kbList != null) {
            fragmentList = new ArrayList<KnowledgeBaseDocument>();
            for (KnowledgeBase kb : kbList) {
                if (kb != null) {
                    if (kb.getFragmentList() != null) {
                        fragmentList.addAll(kb.getKBDocs());
                    }
                }
            }
        }
        return fragmentList;
    }

    /**
     * Returns a list of Elements with Element.ElementType = ENTITY i.e. entities
     * contained in KBTrie
     * @return list of Elements with Element.ElementType = ENTITY
     */
    public List<Element> getEntityList() {
        List<Element> entityList = null;
        if (kbList != null) {
            entityList = new ArrayList<Element>();
            for (KnowledgeBase kb : kbList) {
                if (kb != null) {
                    List<Element> entityList1 = kb.getEntityList();
                    if (entityList1 != null) {
                        entityList.addAll(entityList1);
                    }
                }
            }
        }
        return entityList;
    }

    /**
     * Groups entities contained in the KBTrie by type (ElementOntologySet)
     * @return Map containing <ElementOntologySet, List<Element>>
     */
    public Map<ElementOntologySet, List<Element>> getEntityList_Map() {
        Map<ElementOntologySet, List<Element>> entityList = null;
        if (kbList != null) {
            for (KnowledgeBase kb : kbList) {
                if (kb != null) {
                    entityList = new HashMap<ElementOntologySet, List<Element>>();
                    if (kb.getEntityList() != null) {
                        for (Element entity : kb.getEntityList()) {
                            ElementOntologySet branchID = entity
                                    .getOntologies();
                            List<Element> eList = entityList.get(branchID);
                            if (eList == null) {
                                eList = new ArrayList<Element>();
                                entityList.put(branchID, eList);
                            }
                            eList.add(entity);
                        }
                    }
                }
            }
        }
        return entityList;
    }

    /**
     * Groups entity values contained in the KBTrie by type (ElementOntologySet)
     * @return Map containing <ElementOntologySet, List<String>>
     */
    public Map<ElementOntologySet, Set<String>> getEntityValues() {
        return getEntityValues(getEntityList());
    }

    /**
     * 
     * @param entityList
     * @return 
     */
    public static Map<ElementOntologySet, Set<String>> getEntityValues(
            List<Element> entityList) {
        if (entityList != null) {
            Map<ElementOntologySet, Set<String>> entityValues = new HashMap<ElementOntologySet, Set<String>>();
            for (Element entity : entityList) {
                if (entity == null) 
                    continue;
                ElementOntologySet entityTypeID = entity.getOntologies();
                Set<String> charList = entityValues
                        .get(entityTypeID) == null ? new HashSet<String>()
                        : entityValues.get(entityTypeID);
                String eVal = entity.getValue();

                if(eVal==null)
                    continue;
                charList.add(eVal);
                entityValues.put(entityTypeID, charList);
                
            }
            return entityValues;
        }
        return null;
    }

    /**
     * Returns the regular expression generated using all elements with the same 
     * ElementOntologySet as e excluding e
     * @param e <code>Element</code> to be excluded from regular expression generation
     * @param entityList list of entities
     * @return 
     */
    public Map<ElementOntologySet, Pattern> getEntityPatterns(Element e,
            List<Element> entityList) {
        Map<ElementOntologySet, Pattern> entityTypIDRegex = null;
        if (e != null) {
            if (regexWithoutEntity == null || regexWithoutEntity.get(e) == null) {
                entityTypIDRegex = getEntityTypeIDRegex(getEntityValues(entityList));
                if (entityTypIDRegex != null) {
                    regexWithoutEntity = regexWithoutEntity == null ? new HashMap<Element, Map<ElementOntologySet, Pattern>>()
                            : regexWithoutEntity;
                    regexWithoutEntity.put(e, entityTypIDRegex);
                }
            }
        }
        return entityTypIDRegex;
    }

    public static Map<ElementOntologySet, Pattern> getEntityTypeIDRegex(
            Map<ElementOntologySet, Set<String>> entityValues) {
        if (entityValues != null) {
            Map<ElementOntologySet, Pattern> regexMap = new HashMap<ElementOntologySet, Pattern>();
            DynamicRegexGenerator dynamicRegexGenerator = new DynamicRegexGenerator(threshold);
            Iterator<ElementOntologySet> iter = entityValues.keySet()
                    .iterator();
            while (iter.hasNext()) {
                ElementOntologySet entityTypeID = iter.next();
                Pattern pattern = dynamicRegexGenerator.defineRegex(
                        entityValues.get(entityTypeID));
                if (pattern != null) {
                    regexMap.put(entityTypeID, pattern);
                }
            }
            return regexMap;
        }
        return null;
    }

    /**
     * 
     * @param regexMap
     * @return 
     */
    private RegexCompare generatePatternComparisons(
            Map<ElementOntologySet, Pattern> regexMap) {
        if (regexMap != null) {
            RegexCompare<ElementOntologySet> regCompare = new RegexCompare();
            regCompare.setRegex(regexMap);
            regCompare.generatePatternComparisons();
//            System.out.println("Finished RegexCompare.generatePatternComparisons()");
            return regCompare;
        }
        return null;
    }

    /**
     * 
     */
    private void generatePatternComparisons() {
        allEntityTypeIDRegex = getEntityTypeIDRegex(getEntityValues());
        patternComparison = generatePatternComparisons(allEntityTypeIDRegex);
    }

    /**
     * Normalizes the KBTrie
     * Normalization of KBTrie involves 
     * 1) Generating regular expression for each entity type (ElementOntologySet)
     * 2) Normalizing left and right context tries
     *      count total occurrence of all left and right context in KB
     *      append ambiguous contexts
     */
    public void normalizeModel() {
        generatePatternComparisons();
        this.LEFT_CONTEXT_TRIE.normalize(kbList);
        this.RIGHT_CONTEXT_TRIE.normalize(kbList);        normalize = true;
    }

    /**
     * 
     * @return ElementOntologyFactory
     */
    public ElementOntologyFactory getElementOntologyFactory() {
        return elementOntologyFactory;
    }

    public ElementOntologySetGraph getElementOntologySetGraph() {
        return elementOntologySetGraph;
    }

    public RegexCompare getPatternComparison() {
        return patternComparison;
    }
    
    
    


}
